{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import sys\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# Add the project root to Python path\n",
    "project_root = Path(\"/Users/salahalzubi/cursor_projects/SentientResearchAgent\")\n",
    "sys.path.insert(0, str(project_root / \"src\"))\n",
    "\n",
    "# Basic imports to test the fix\n",
    "from agno.agent import Agent as AgnoAgent\n",
    "from agno.models.litellm import LiteLLM\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import List, Optional\n",
    "import json\n",
    "import re\n",
    "\n",
    "# Define the models directly to avoid circular imports\n",
    "class SubTask(BaseModel):\n",
    "    goal: str = Field(..., description=\"Precise description of the sub-task goal.\")\n",
    "    task_type: str = Field(..., description=\"Type of task (e.g., 'WRITE', 'THINK', 'SEARCH').\")\n",
    "    node_type: str = Field(..., description=\"Node type ('EXECUTE' for atomic, 'PLAN' for complex).\")\n",
    "    depends_on_indices: Optional[List[int]] = Field(default_factory=list)\n",
    "\n",
    "class PlanOutput(BaseModel):\n",
    "    sub_tasks: List[SubTask] = Field(..., description=\"List of planned sub-tasks.\")\n",
    "\n",
    "# Your JSON extraction function (from the fix)\n",
    "def extract_json_block_with_backticks(raw_response: str) -> Optional[str]:\n",
    "    \"\"\"Extract JSON content from LLM response that may contain thinking tags and backticks.\"\"\"\n",
    "    \n",
    "    backtick_patterns = [\n",
    "        r'```(?:json)?\\s*\\n([\\s\\S]*?)\\n?```',  # Extract content inside backticks\n",
    "    ]\n",
    "    \n",
    "    for pattern in backtick_patterns:\n",
    "        matches = re.findall(pattern, raw_response, re.IGNORECASE)\n",
    "        for match in matches:\n",
    "            cleaned_match = match.strip()\n",
    "            if cleaned_match.startswith('{') or cleaned_match.startswith('['):\n",
    "                try:\n",
    "                    json.loads(cleaned_match)  # Validate JSON\n",
    "                    print(f\"✅ Found valid JSON code block: {len(cleaned_match)} chars\")\n",
    "                    return cleaned_match\n",
    "                except json.JSONDecodeError:\n",
    "                    continue\n",
    "\n",
    "    print(\"❌ No valid JSON block found\")\n",
    "    return None\n",
    "\n",
    "DEEP_RESEARCH_PLANNER_SYSTEM_MESSAGE = \"\"\"You are an elite Hierarchical Planning Agent. Your sole purpose is to receive a complex question or research goal and decompose it into a precise, logical, and actionable sequence of sub-tasks. You specialize in planning for information-retrieval, reasoning, and synthesis tasks. You operate with surgical precision, ensuring every plan is coherent, efficient, and directly aimed at producing a complete and accurate answer. You do not execute tasks; you only create the plan.\n",
    "\n",
    "**Input Schema:**\n",
    "\n",
    "You will receive input in JSON format with the following fields:\n",
    "\n",
    "*   `current_task_goal` (string, mandatory): The specific goal for this planning instance.\n",
    "*   `overall_objective` (string, mandatory): The ultimate high-level goal of the entire operation. This helps maintain alignment.\n",
    "*   `parent_task_goal` (string, optional): The goal of the immediate parent task that led to this decomposition. Null if this is the root task.\n",
    "*   `planning_depth` (integer, optional): Current recursion depth (e.g., 0 for initial, 1 for sub-tasks).\n",
    "*   `execution_history_and_context` (object, mandatory):\n",
    "    *   `prior_sibling_task_outputs` (array of objects, optional): Outputs from tasks at the same hierarchical level that executed before this planning step. Each object contains:\n",
    "        *   `task_goal` (string): Goal of the sibling task.\n",
    "        *   `outcome_summary` (string): Brief summary of what the sibling task achieved or produced.\n",
    "        *   `full_output_reference_id` (string, optional): ID to fetch the full output if needed.\n",
    "    *   `relevant_ancestor_outputs` (array of objects, optional): Key outputs from parent or higher-level tasks crucial for `current_task_goal`. Each object similar to sibling outputs.\n",
    "    *   `global_knowledge_base_summary` (string, optional): Brief summary/keywords of available global knowledge.\n",
    "*   `replan_request_details` (object, optional): If this is a re-plan, this object contains structured feedback. Null otherwise.\n",
    "    *   `failed_sub_goal` (string): The specific sub-goal related to `current_task_goal` that previously failed.\n",
    "    *   `reason_for_failure_or_replan` (string): Detailed explanation of the failure or re-plan need.\n",
    "    *   `previous_attempt_output_summary` (string, optional): Summary of the failed attempt's output.\n",
    "    *   `specific_guidance_for_replan` (string, optional): Concrete suggestions for the re-plan.\n",
    "*   `global_constraints_or_preferences` (array of strings, optional): E.g., \"Prioritize accuracy\", \"Maximum 3 sub-tasks\".\n",
    "\n",
    "**Core Task:**\n",
    "\n",
    "1.  Analyze the `current_task_goal` in the context of `overall_objective`, `parent_task_goal`, and available `execution_history_and_context`.\n",
    "2.  Decompose `current_task_goal` into a list of **3 to 6 granular sub-tasks.** If a goal is exceptionally complex and absolutely requires more than 6 sub-tasks to maintain clarity and avoid overly broad steps, you may slightly exceed this, but strive for conciseness. Aim for sub-tasks that represent meaningful, coherent units of work. While `EXECUTE` tasks should be specific, avoid breaking down a goal into excessively small pieces if a slightly larger, but still focused and directly actionable, `EXECUTE` task is feasible for a specialized agent. Prioritize clarity and manageability over maximum possible decomposition.\n",
    "3.  For each sub-task, define:\n",
    "    *   `goal` (string): The specific goal. Ensure sub-task goals are distinct and avoid significant overlap with sibling tasks in the current plan.\n",
    "    *   `task_type` (string): 'WRITE', 'THINK', or 'SEARCH'.\n",
    "    *   `node_type` (string): 'EXECUTE' (atomic) or 'PLAN' (needs more planning).\n",
    "    *   `depends_on_indices` (list of integers, optional): A list of 0-based indices of other sub-tasks *in the current list of sub-tasks you are generating* that this specific sub-task directly depends on. Example: If sub-task at index 2 depends on sub-task at index 0 and sub-task at index 1, this field would be `[0, 1]`. If a sub-task can start as soon as the parent plan is approved (i.e., it doesn't depend on any other sibling sub-tasks in *this* plan), this should be an empty list `[]`. Use this to define sequential dependencies when one sub-task in your plan needs the output of another sub-task from the *same* plan. Ensure indices are valid and refer to previously listed sub-tasks in your current plan.\n",
    "4.  **Task Ordering and Dependencies**:\n",
    "    *   List sub-tasks in a logical order.\n",
    "    *   Use `depends_on_indices` to explicitly state if a sub-task requires the completion of one or more *other sub-tasks from the current plan* before it can start.\n",
    "    *   If tasks are largely independent and can run in parallel, their `depends_on_indices` should be `[]`.\n",
    "\n",
    "**Re-planning Logic**: \n",
    "\n",
    "If `replan_request_details` is provided:\n",
    "    *   Pay **critical attention** to `reason_for_failure_or_replan` and `specific_guidance_for_replan`.\n",
    "    *   Your new plan **MUST** address the failure by:\n",
    "        *   Being more granular for the `failed_sub_goal`.\n",
    "        *   Altering the approach (e.g., different `task_type`s).\n",
    "        *   Suggesting different information gathering if context was missing.\n",
    "        *   Modifying sub-task goals based on `specific_guidance_for_replan`.\n",
    "        *   Adjusting `depends_on_indices` if the previous dependency structure was flawed.\n",
    "    *   Ensure the new plan for `current_task_goal` explicitly mitigates the previous failure.\n",
    "\n",
    "**Planning Tips (Leveraging New Input):**\n",
    "\n",
    "1.  **Context is Key**: Use `prior_sibling_task_outputs` to build sequentially (if logically dependent) and avoid redundancy. Leverage `relevant_ancestor_outputs`.\n",
    "2.  **Mutual Exclusivity & Complementation**:\n",
    "    *   Strive for sub-tasks that cover different aspects of the `current_task_goal` without significant overlap. They should be complementary, together achieving the parent goal.\n",
    "    *   Before finalizing sub-tasks, review them as a set: Do they make sense together? Is there redundancy? Are there gaps? Are dependencies correctly defined using `depends_on_indices`?\n",
    "3.  **CRITICAL - Balanced Granularity for SEARCH Tasks**:\n",
    "    *   **`SEARCH/EXECUTE` Specificity**: A `SEARCH/EXECUTE` sub-task goal **MUST** be so specific that it typically targets a single fact, statistic, definition, or a very narrow aspect of a topic.\n",
    "        *   *Good `SEARCH/EXECUTE` examples*: \"Find the 2023 import tariff rate for Chinese-made solar panels in the US.\", \"List the main arguments for the Jones Act.\"\n",
    "        *   *Bad `SEARCH/EXECUTE` examples (these should be `SEARCH/PLAN` or broken down)*: \"Research US solar panel tariffs.\", \"Understand the Jones Act.\"\n",
    "    *   **Avoiding Over-Fragmentation**: While specificity is key, if multiple *very small, extremely closely related pieces of data* can be retrieved with a single, well-crafted, targeted search query (and an agent can easily parse them), you can group them into one `SEARCH/EXECUTE` task. Example: Instead of three tasks \"Find 2022 EV sales\", \"Find 2023 EV sales\", \"Find 2024 EV sales\", one task \"Find annual US EV sales figures for 2022, 2023, and 2024\" is acceptable if the search agent can handle it. However, do not combine distinct conceptual questions.\n",
    "    *   **When to use `SEARCH/PLAN`**: If a research sub-goal still requires investigating multiple *distinct conceptual areas* or is too broad for one or two highly targeted queries (even if slightly grouped as above), that sub-task **MUST** be `task_type: 'SEARCH'` and `node_type: 'PLAN'`. This ensures it gets further decomposed.\n",
    "\n",
    "**Required Output Attributes per Sub-Task:**\n",
    "`goal`, `task_type` (string: 'WRITE', 'THINK', or 'SEARCH'), `node_type` (string: 'EXECUTE' or 'PLAN'), `depends_on_indices` (list of integers).\n",
    "\n",
    "**Output Format:**\n",
    "- Respond ONLY with a JSON list of sub-task objects.\n",
    "- Or an empty list if the `current_task_goal` cannot or should not be broken down further (e.g., it's already atomic enough given the context).\n",
    "\n",
    "---\n",
    "### Examples\n",
    "\n",
    "[BEGIN]\n",
    "**Input:**\n",
    "```json\n",
    "{\n",
    "  \"current_task_goal\": \"Explain how the invention of the transistor led to the development of the modern internet.\",\n",
    "  \"overall_objective\": \"Answer a user's question about technological history.\",\n",
    "  \"execution_history_and_context\": {}\n",
    "}\n",
    "\n",
    "\n",
    "**Output:**\n",
    "[\n",
    "  {\n",
    "    \"goal\": \"Find the date and primary function of the transistor's invention, focusing on its role in replacing vacuum tubes.\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Research how transistors enabled the creation of smaller, more reliable, and more powerful computers via integrated circuits (microchips).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [0]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Research the origins of ARPANET and identify its core requirement for a network of interconnected, powerful computers at various nodes.\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [1]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Synthesize the findings to construct the causal chain: transistors led to powerful/small computers (via ICs), which were a necessary precondition for a distributed network like ARPANET, the precursor to the internet.\",\n",
    "    \"task_type\": \"THINK\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [2]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Write a clear, step-by-step explanation answering the original question.\",\n",
    "    \"task_type\": \"WRITE\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [3]\n",
    "  }\n",
    "]\n",
    "\n",
    "[END]\n",
    "\n",
    "[BEGIN]\n",
    "**Input:**\n",
    "{\n",
    "  \"current_task_goal\": \"Compare and contrast the economic policies of Reaganomics in the 1980s and Clintonomics in the 1990s, focusing on their stated goals and impact on the US national debt.\",\n",
    "  \"overall_objective\": \"Answer a user's question about economic policy.\",\n",
    "  \"execution_history_and_context\": {}\n",
    "}\n",
    "\n",
    "**Output:**\n",
    "[\n",
    "  {\n",
    "    \"goal\": \"Identify the core principles and stated goals of Reaganomics (e.g., supply-side economics, tax cuts, deregulation).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Identify the core principles and stated goals of Clintonomics (e.g., deficit reduction, targeted investments, trade liberalization).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Find US national debt figures for the periods 1981-1989 and 1993-2001.\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Analyze and summarize the similarities and differences in the stated goals and principles of the two economic policies.\",\n",
    "    \"task_type\": \"THINK\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [0, 1]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Synthesize the policy principles and debt figures to compare the actual impact of each administration's policies on the national debt.\",\n",
    "    \"task_type\": \"THINK\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [0, 1, 2]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Write a final answer that first compares the policies' goals and then contrasts their effects on the national debt, citing the data found.\",\n",
    "    \"task_type\": \"WRITE\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [3, 4]\n",
    "  }\n",
    "]\n",
    "\n",
    "[END]\n",
    "\n",
    "[BEGIN]\n",
    "**Input:**\n",
    "\n",
    "{\n",
    "  \"current_task_goal\": \"What is Quantum Computing, and what are its most significant potential benefits and risks?\",\n",
    "  \"overall_objective\": \"Provide a comprehensive but accessible explanation of a complex topic.\",\n",
    "  \"execution_history_and_context\": {}\n",
    "}\n",
    "\n",
    "**Output:**\n",
    "[\n",
    "  {\n",
    "    \"goal\": \"Find a clear, concise definition of quantum computing, including its core principles like superposition and entanglement.\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Identify 3-4 of the most significant potential benefits and applications of quantum computing (e.g., drug discovery, financial modeling, materials science).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Identify 3-4 of the most significant risks or challenges associated with quantum computing (e.g., breaking current encryption, high error rates, decoherence).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Synthesize the collected information to structure a balanced answer: first the definition, then the benefits, and finally the risks.\",\n",
    "    \"task_type\": \"THINK\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [0, 1, 2]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Write the final explanation in clear, accessible language, suitable for a non-expert audience.\",\n",
    "    \"task_type\": \"WRITE\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [3]\n",
    "  }\n",
    "]\n",
    "\n",
    "[END]\n",
    "\n",
    "[BEGIN]\n",
    "**Input:**\n",
    "\n",
    "{\n",
    "  \"current_task_goal\": \"What are the primary challenges and proposed solutions for establishing a sustainable human colony on Mars?\",\n",
    "  \"overall_objective\": \"Answer a user's question about space colonization.\",\n",
    "  \"execution_history_and_context\": {}\n",
    "}\n",
    "\n",
    "**Output:**\n",
    "[\n",
    "  {\n",
    "    \"goal\": \"Identify the top 3-4 primary survival challenges for a Mars colony (e.g., radiation, thin atmosphere/pressure, resource scarcity, psychological effects).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": []\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"For each identified challenge, research the leading proposed solutions (e.g., for radiation: subsurface habitats, magnetic shielding; for resources: In-Situ Resource Utilization (ISRU) for water and oxygen).\",\n",
    "    \"task_type\": \"SEARCH\",\n",
    "    \"node_type\": \"PLAN\",\n",
    "    \"depends_on_indices\": [0]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Synthesize the research by mapping each challenge directly to its most promising proposed solution(s).\",\n",
    "    \"task_type\": \"THINK\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [1]\n",
    "  },\n",
    "  {\n",
    "    \"goal\": \"Write a structured answer that first lists the primary challenges and then, for each challenge, explains the corresponding proposed solutions.\",\n",
    "    \"task_type\": \"WRITE\",\n",
    "    \"node_type\": \"EXECUTE\",\n",
    "    \"depends_on_indices\": [2]\n",
    "  }\n",
    "]\n",
    "[END]\n",
    "\"\"\" \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enhanced extraction function that handles more cases\n",
    "def extract_json_from_response(raw_response: str) -> Optional[str]:\n",
    "    \"\"\"\n",
    "    Enhanced JSON extraction that handles multiple formats.\n",
    "    \"\"\"\n",
    "    print(f\"🔍 Analyzing response (length: {len(raw_response)})\")\n",
    "    \n",
    "    # Method 1: JSON in triple backticks\n",
    "    backtick_patterns = [\n",
    "        r'```(?:json)?\\s*\\n([\\s\\S]*?)\\n?```',  # Standard backticks\n",
    "        r'```([\\s\\S]*?)```',  # Simple backticks\n",
    "    ]\n",
    "    \n",
    "    for pattern in backtick_patterns:\n",
    "        matches = re.findall(pattern, raw_response, re.IGNORECASE)\n",
    "        for match in matches:\n",
    "            cleaned = match.strip()\n",
    "            if cleaned.startswith(('[', '{')):\n",
    "                try:\n",
    "                    json.loads(cleaned)\n",
    "                    print(f\"✅ Found JSON in backticks: {len(cleaned)} chars\")\n",
    "                    return cleaned\n",
    "                except json.JSONDecodeError:\n",
    "                    continue\n",
    "    \n",
    "    # Method 2: Remove <think> tags and look for JSON arrays/objects\n",
    "    # Remove thinking tags\n",
    "    cleaned_response = re.sub(r'<think>.*?</think>', '', raw_response, flags=re.DOTALL | re.IGNORECASE)\n",
    "    print(f\"🧹 After removing <think> tags: {len(cleaned_response)} chars\")\n",
    "    \n",
    "    # Look for JSON arrays or objects\n",
    "    json_patterns = [\n",
    "        r'(\\[[\\s\\S]*?\\])',  # JSON arrays\n",
    "        r'(\\{[\\s\\S]*?\\})',  # JSON objects  \n",
    "    ]\n",
    "    \n",
    "    for pattern in json_patterns:\n",
    "        matches = re.findall(pattern, cleaned_response)\n",
    "        for match in matches:\n",
    "            try:\n",
    "                parsed = json.loads(match)\n",
    "                # Validate it looks like a task list\n",
    "                if isinstance(parsed, list) and len(parsed) > 0:\n",
    "                    if isinstance(parsed[0], dict) and 'goal' in parsed[0]:\n",
    "                        print(f\"✅ Found JSON array without backticks: {len(match)} chars\")\n",
    "                        return match\n",
    "                elif isinstance(parsed, dict) and 'sub_tasks' in parsed:\n",
    "                    print(f\"✅ Found JSON object without backticks: {len(match)} chars\")\n",
    "                    return match\n",
    "            except json.JSONDecodeError:\n",
    "                continue\n",
    "    \n",
    "    # Method 3: Look for JSON-like structures line by line\n",
    "    lines = cleaned_response.split('\\n')\n",
    "    json_lines = []\n",
    "    in_json = False\n",
    "    bracket_count = 0\n",
    "    \n",
    "    for line in lines:\n",
    "        line = line.strip()\n",
    "        if line.startswith('[') or line.startswith('{'):\n",
    "            in_json = True\n",
    "            json_lines = [line]\n",
    "            bracket_count = line.count('[') + line.count('{') - line.count(']') - line.count('}')\n",
    "        elif in_json:\n",
    "            json_lines.append(line)\n",
    "            bracket_count += line.count('[') + line.count('{') - line.count(']') - line.count('}')\n",
    "            if bracket_count <= 0:\n",
    "                # Try to parse the accumulated JSON\n",
    "                potential_json = '\\n'.join(json_lines)\n",
    "                try:\n",
    "                    parsed = json.loads(potential_json)\n",
    "                    print(f\"✅ Found JSON by line parsing: {len(potential_json)} chars\")\n",
    "                    return potential_json\n",
    "                except json.JSONDecodeError:\n",
    "                    pass\n",
    "                in_json = False\n",
    "                json_lines = []\n",
    "                bracket_count = 0\n",
    "    \n",
    "    print(\"❌ No valid JSON found with any method\")\n",
    "    return None\n",
    "\n",
    "# Enhanced test function with better debugging\n",
    "async def debug_deepresearch_response():\n",
    "    \"\"\"Debug the actual response from DeepResearchPlanner.\"\"\"\n",
    "    \n",
    "    print(\"🔍 DEBUGGING DeepResearchPlanner Response\")\n",
    "    print(\"=\"*60)\n",
    "    \n",
    "    model = LiteLLM(id=\"fireworks_ai/accounts/fireworks/models/deepseek-r1-0528\")\n",
    "    \n",
    "    # Test WITHOUT response_model first to see raw output\n",
    "    agent_raw = AgnoAgent(\n",
    "        model=model,\n",
    "        system_message=DEEP_RESEARCH_PLANNER_SYSTEM_MESSAGE,\n",
    "        name=\"DeepResearchPlanner_Debug\",\n",
    "        # No response_model - get raw output\n",
    "        markdown=False\n",
    "    )\n",
    "    \n",
    "    test_prompt = \"\"\"Overall Objective: Research machine learning in healthcare\n",
    "\n",
    "Current Task Goal: Create a 3-task research plan to analyze ML applications in diagnostics, treatment planning, and drug discovery\n",
    "\n",
    "Based on the 'Current Task Goal', generate a plan to achieve it.\"\"\"\n",
    "    \n",
    "    try:\n",
    "        print(\"📤 Getting raw response from model...\")\n",
    "        result = await agent_raw.arun(test_prompt)\n",
    "        \n",
    "        if hasattr(result, 'content'):\n",
    "            content = result.content\n",
    "            if asyncio.iscoroutine(content):\n",
    "                content = await content\n",
    "        else:\n",
    "            content = result\n",
    "        \n",
    "        print(f\"📄 Full raw response:\")\n",
    "        print(\"-\" * 50)\n",
    "        print(content)\n",
    "        print(\"-\" * 50)\n",
    "        \n",
    "        # Try enhanced extraction\n",
    "        print(f\"\\n🔧 Trying enhanced JSON extraction...\")\n",
    "        extracted = extract_json_from_response(content)\n",
    "        \n",
    "        if extracted:\n",
    "            print(f\"\\n✅ Extracted JSON:\")\n",
    "            print(extracted)\n",
    "            \n",
    "            # Try to convert to PlanOutput\n",
    "            try:\n",
    "                parsed_data = json.loads(extracted)\n",
    "                if isinstance(parsed_data, list):\n",
    "                    sub_tasks = [SubTask(**item) for item in parsed_data]\n",
    "                    plan = PlanOutput(sub_tasks=sub_tasks)\n",
    "                    print(f\"\\n🎉 SUCCESS: Created PlanOutput with {len(plan.sub_tasks)} tasks\")\n",
    "                    return plan\n",
    "                elif isinstance(parsed_data, dict) and 'sub_tasks' in parsed_data:\n",
    "                    plan = PlanOutput(**parsed_data)\n",
    "                    print(f\"\\n🎉 SUCCESS: Created PlanOutput from dict\")\n",
    "                    return plan\n",
    "            except Exception as e:\n",
    "                print(f\"\\n❌ Failed to create PlanOutput: {e}\")\n",
    "        else:\n",
    "            print(f\"\\n❌ No JSON could be extracted\")\n",
    "            \n",
    "            # Let's see if we can find any JSON-like patterns\n",
    "            print(f\"\\n🔍 Looking for any JSON-like patterns...\")\n",
    "            \n",
    "            # Look for common JSON indicators\n",
    "            if '[' in content and ']' in content:\n",
    "                print(f\"   Found array brackets\")\n",
    "                start = content.find('[')\n",
    "                end = content.rfind(']') + 1\n",
    "                if start < end:\n",
    "                    potential = content[start:end]\n",
    "                    print(f\"   Potential array: {potential[:200]}...\")\n",
    "            \n",
    "            if '{' in content and '}' in content:\n",
    "                print(f\"   Found object brackets\")\n",
    "                start = content.find('{')\n",
    "                end = content.rfind('}') + 1\n",
    "                if start < end:\n",
    "                    potential = content[start:end]\n",
    "                    print(f\"   Potential object: {potential[:200]}...\")\n",
    "        \n",
    "        return None\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Debug failed: {e}\")\n",
    "        import traceback\n",
    "        traceback.print_exc()\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔍 DEBUGGING DeepResearchPlanner Response\n",
      "============================================================\n",
      "📤 Getting raw response from model...\n",
      "📄 Full raw response:\n",
      "--------------------------------------------------\n",
      "<think>\n",
      "We are given an overall objective and a current task goal that explicitly asks for a research plan focused on three areas.\n",
      " The task goal is: \"Create a 3-task research plan to analyze ML applications in diagnostics, treatment planning, and drug discovery\"\n",
      " This is a planning task (PLAN node) and the task_type is THINK because we are creating a plan, not directly searching for information.\n",
      " However, note that the current task goal is to \"Create\" a plan. Therefore, we are generating a plan for a plan? \n",
      " But let's clarify: the current task is to create a research plan that has 3 tasks (each covering one area). \n",
      " Since the output of this task will be a plan (a list of tasks) for researching the three areas, we are acting in a planning capacity.\n",
      "\n",
      " However, note the role: We are a Hierarchical Planning Agent. We are to decompose the current task goal.\n",
      "\n",
      " The current task goal is to create a 3-task research plan. How do we decompose this?\n",
      " We can break it down into three sub-tasks, each focusing on one area (diagnostics, treatment planning, drug discovery) and then a synthesis step?\n",
      "\n",
      " But wait: the goal is to \"create a 3-task research plan\". So the output should be a plan that has three tasks? \n",
      " Actually, the current task goal is a request for a plan that will have three tasks.\n",
      "\n",
      " However, our system is designed to break down the current task goal into sub-tasks that we then assign as either EXECUTE or PLAN.\n",
      "\n",
      " Since the current task goal is to \"create\" a plan, we can interpret it as:\n",
      "\n",
      " Step 1: Plan the task for diagnostics.\n",
      " Step 2: Plan the task for treatment planning.\n",
      " Step 3: Plan the task for drug discovery.\n",
      "\n",
      " But note: the output we are to produce is a list of sub-tasks that when executed will result in the creation of the 3-task research plan.\n",
      "\n",
      " However, the current task goal is abstract: we are not doing the research, we are creating a research plan.\n",
      "\n",
      " Let's model:\n",
      "\n",
      " We are to create a research plan that covers the three areas. Each of the three tasks in the research plan will be a separate research assignment.\n",
      "\n",
      " But note: the current task goal is to create the plan. Therefore, we are generating the structure of the research.\n",
      "\n",
      " How to break down the current task goal?\n",
      "\n",
      " Since the goal is to create a 3-task research plan, we can break it down into:\n",
      "\n",
      " 1. Define the scope and research questions for the diagnostics application area.\n",
      " 2. Define the scope and research questions for the treatment planning application area.\n",
      " 3. Define the scope and research questions for the drug discovery application area.\n",
      "\n",
      " Then, we can have a step to compile the three tasks into a coherent plan.\n",
      "\n",
      " However, note that the current task goal says \"3-task research plan\", so we must output exactly three tasks? \n",
      " Actually, the goal is to create a plan that has three tasks. So we are creating three tasks? \n",
      "\n",
      " But wait: the current task goal is to \"create\" the plan. So we are going to output a plan that has three tasks? \n",
      "\n",
      " However, the instructions for the Hierarchical Planning Agent are that we break down the current task goal into sub-tasks.\n",
      "\n",
      " We are at a planning node. We are to break down the task of \"creating a 3-task research plan\" into steps.\n",
      "\n",
      " We propose:\n",
      "\n",
      " 1. For the 'diagnostics' area: Define a research task that will cover machine learning applications in diagnostics. This is a planning task because we are defining the task (which will later be executed by another agent). So the output of step 1 is a sub-plan for the diagnostics research.\n",
      "\n",
      " 2. Similarly for treatment planning.\n",
      "\n",
      " 3. Similarly for drug discovery.\n",
      "\n",
      " 4. Then, we need to compile these three research tasks into a single plan document? \n",
      "\n",
      " But note: the current task goal is to create the plan. So the plan we are creating is the plan that has three tasks. \n",
      "\n",
      " However, the problem is: we are a planning agent. Our job is to break down the current task goal. We are not executing the research. We are creating the plan for the research.\n",
      "\n",
      " Therefore, the output of our planning (the current decomposition) should be the steps necessary to create the three research tasks.\n",
      "\n",
      " How about:\n",
      "\n",
      " We break down the current task goal into:\n",
      "\n",
      "  - Sub-task 1: Formulate the research task for diagnostics (what exactly to research? what questions to answer? what kind of information to find?). This is a THINK/PLAN task? But note: we are creating a plan for the research, so we are defining the parameters of the research task.\n",
      "\n",
      "  - Sub-task 2: Formulate the research task for treatment planning.\n",
      "\n",
      "  - Sub-task 3: Formulate the research task for drug discovery.\n",
      "\n",
      "  - Sub-task 4: Combine these three task definitions into a structured research plan (a list of tasks).\n",
      "\n",
      " But the goal is a 3-task research plan. So step 4 would be to output the plan as an object with three tasks.\n",
      "\n",
      " However, note that the current task goal does not require compiling beyond having the three tasks defined? \n",
      "\n",
      " Alternatively, we can see the current task goal as:\n",
      "\n",
      "   The output of this planning step should be the three research tasks (as sub-tasks that are PLAN nodes?).\n",
      "\n",
      " But wait: the current task goal is to \"Create a 3-task research plan\". So the output of this entire planning step we are doing now is the three tasks? \n",
      "\n",
      " However, the instructions say: we are to break down the current task goal into sub-tasks. And the output is a list of sub-tasks.\n",
      "\n",
      " Therefore, we must break down the task of creating the 3-task research plan into steps.\n",
      "\n",
      " But note: we are the planning agent. We are doing the decomposition. We are not executing the research. We are creating the plan for the research.\n",
      "\n",
      " We can break down the task of creating the research plan into:\n",
      "\n",
      "  Step 1: Research the key aspects of ML in diagnostics to define the scope and specific questions for the diagnostics research task. (This is a SEARCH/EXECUTE because we need to understand what to include in the diagnostics research task? But wait, we are creating the plan, not doing the research. We are just designing the research task.)\n",
      "\n",
      "  Step 2: Similarly for treatment planning.\n",
      "\n",
      "  Step 3: Similarly for drug discovery.\n",
      "\n",
      "  Step 4: Write the research plan document that lists the three tasks.\n",
      "\n",
      " However, that would be 4 steps. And the goal is a 3-task research plan. \n",
      "\n",
      " Alternatively, we can view it as:\n",
      "\n",
      "   We are going to think about each of the three areas and define the research task for each.\n",
      "\n",
      "   Then we output the three tasks as the plan.\n",
      "\n",
      " But note: our role is to output a list of sub-tasks that will achieve the current task goal.\n",
      "\n",
      " So the sub-tasks we output are:\n",
      "\n",
      "   [\n",
      "      {\n",
      "        goal: \"Define the research task for machine learning in diagnostics, including the scope and key questions to be answered.\",\n",
      "        task_type: \"THINK\",   // because we are formulating a task\n",
      "        node_type: \"EXECUTE\", // because we can do it in one step? Or PLAN? \n",
      "        depends_on_indices: []\n",
      "      },\n",
      "      ... similarly for treatment planning and drug discovery ...\n",
      "   ]\n",
      "\n",
      " Then, after these three, we don't need a fourth? Because the current task goal is to create the plan, and after defining the three tasks, we have the plan.\n",
      "\n",
      " But note: the output of our decomposition must be such that when these sub-tasks are executed, they result in the creation of the 3-task research plan.\n",
      "\n",
      " However, the Hierarchical Planning Agent does not execute the tasks. We are only creating the plan.\n",
      "\n",
      " Therefore, the plan we create (the output of this step) is the list of three tasks (each being a research task). \n",
      "\n",
      " How do we represent this? \n",
      "\n",
      " We are to break down the current task goal (which is to create the plan) into steps that will produce that plan.\n",
      "\n",
      " We can have:\n",
      "\n",
      "   Step 1: Think about what the diagnostics research task should entail.\n",
      "   Step 2: Think about what the treatment planning research task should entail.\n",
      "   Step 3: Think about what the drug discovery research task should entail.\n",
      "\n",
      "   Then, we have a step to actually create the plan document? But the goal is the plan, not a document.\n",
      "\n",
      " Actually, the output of the current task (when executed) is the plan (which is a list of tasks). So the steps we break down must lead to generating that list.\n",
      "\n",
      " However, in our decomposition, we are the planning agent. We are creating the plan for the task \"create a 3-task research plan\". \n",
      "\n",
      " We are at the planning stage. Our output will be a list of sub-tasks that when executed will produce the three tasks.\n",
      "\n",
      " But note: the three tasks we are to create are actually plans themselves? They are research tasks that will be broken down later.\n",
      "\n",
      " Therefore, we can output:\n",
      "\n",
      "   [\n",
      "        {\n",
      "            \"goal\": \"Define the specifics of the research task for ML applications in diagnostics, including key questions and scope.\",\n",
      "            \"task_type\": \"THINK\",\n",
      "            \"node_type\": \"EXECUTE\",\n",
      "            \"depends_on_indices\": []\n",
      "        },\n",
      "        {\n",
      "            \"goal\": \"Define the specifics of the research task for ML applications in treatment planning, including key questions and scope.\",\n",
      "            \"task_type\": \"THINK\",\n",
      "            \"node_type\": \"EXECUTE\",\n",
      "            \"depends_on_indices\": []\n",
      "        },\n",
      "        {\n",
      "            \"goal\": \"Define the specifics of the research task for ML applications in drug discovery, including key questions and scope.\",\n",
      "            \"task_type\": \"THINK\",\n",
      "            \"node_type\": \"EXECUTE\",\n",
      "            \"depends_on_indices\": []\n",
      "        }\n",
      "   ]\n",
      "\n",
      " And then,\n",
      "--------------------------------------------------\n",
      "\n",
      "🔧 Trying enhanced JSON extraction...\n",
      "🔍 Analyzing response (length: 9340)\n",
      "🧹 After removing <think> tags: 9340 chars\n",
      "✅ Found JSON by line parsing: 201 chars\n",
      "\n",
      "✅ Extracted JSON:\n",
      "{\n",
      "\"goal\": \"Define the specifics of the research task for ML applications in drug discovery, including key questions and scope.\",\n",
      "\"task_type\": \"THINK\",\n",
      "\"node_type\": \"EXECUTE\",\n",
      "\"depends_on_indices\": []\n",
      "}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/salahalzubi/cursor_projects/SentientResearchAgent/.venv/lib/python3.12/site-packages/pydantic/main.py:463: UserWarning: Pydantic serializer warnings:\n",
      "  PydanticSerializationUnexpectedValue(Expected 9 fields but got 5: Expected `Message` - serialized value may not be as expected [input_value=Message(content='<think>\\...er_specific_fields=None), input_type=Message])\n",
      "  PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [input_value=Choices(finish_reason='le...r_specific_fields=None)), input_type=Choices])\n",
      "  return self.__pydantic_serializer__.to_python(\n"
     ]
    }
   ],
   "source": [
    "debug_result = await debug_deepresearch_response()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "print(debug_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
